// Copyright (c) .NET Foundation. All rights reserved.
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.

using System;
using System.Collections.Generic;
using System.IO;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using System.Linq;
using GitHubVulnerabilities2v3.Telemetry;
using GitHubVulnerabilities2v3.Configuration;
using GitHubVulnerabilities2v3.Entities;
using Microsoft.Extensions.Logging;
using Newtonsoft.Json;
using NuGet.Services.Cursor;
using NuGet.Services.Entities;
using NuGet.Services.GitHub.Ingest;
using NuGet.Services.Storage;

namespace GitHubVulnerabilities2v3.Extensions
{
    public class BlobStorageVulnerabilityWriter : IVulnerabilityWriter
    {
        private static readonly string BASE_ENTRY_NAME = "base";
        private static readonly string UPDATE_ENTRY_NAME = "update";
        private static readonly string TimeFormat = "yyyy.MM.dd.HH.mm.ss";
        private static readonly string JsonContentType = "application/json";

        private readonly GitHubVulnerabilities2v3Configuration _configuration;
        private readonly ReadWriteCursor<DateTimeOffset> _cursor;
        private readonly ILogger _logger;
        private readonly IStorageFactory _storageFactory;

        private Dictionary<string, List<Advisory>> _vulnerabilityDict;
        private IStorage _storage;
        private ITelemetryService _telemetryService;
        private DateTimeOffset _firstVulnWrittenTimestamp;

        public BlobStorageVulnerabilityWriter(
            IStorageFactory storageFactory,
            GitHubVulnerabilities2v3Configuration configuration,
            ReadWriteCursor<DateTimeOffset> cursor,
            ILogger<BlobStorageVulnerabilityWriter> logger,
            ITelemetryService telemetryService)
        {
            _storageFactory = storageFactory ?? throw new ArgumentNullException(nameof(storageFactory));
            _configuration = configuration ?? throw new ArgumentNullException(nameof(configuration));
            _cursor = cursor ?? throw new ArgumentNullException(nameof(cursor));
            _logger = logger ?? throw new ArgumentNullException(nameof(logger));
            _telemetryService = telemetryService ?? throw new ArgumentNullException(nameof(telemetryService));

            _storage = _storageFactory.Create();
            if (_storage is AzureStorage azureStorage)
            {
                azureStorage.CompressContent = _configuration.GzipFileContent;
            }

            _firstVulnWrittenTimestamp = DateTimeOffset.MinValue;
            _vulnerabilityDict = new Dictionary<string, List<Advisory>>(StringComparer.OrdinalIgnoreCase);
        }

        public async Task FlushAsync(string outputFileName = null)
        {
            var runMode = await DetermineRunMode(_cursor);
            StringBuilder sb = new StringBuilder();

            using (StringWriter sw = new StringWriter(sb))
            using (JsonWriter textWriter = new JsonTextWriter(sw))
            {
                var serializer = new JsonSerializer();
                serializer.Serialize(textWriter, _vulnerabilityDict);
                await textWriter.FlushAsync();
            }

            var stringContentOutput = sb.ToString();

            int vulnerabilityCount = 0;
            foreach(var vulnerability in _vulnerabilityDict.Values)
            {
                vulnerabilityCount += vulnerability.Count;
            }

            var currentTime = DateTime.UtcNow.ToString(TimeFormat);
            var indexStorageUri = _storage.ResolveUri(_configuration.IndexFileName);
            StringBuilder updateUriPathBuilder = new StringBuilder();
            updateUriPathBuilder.Append(_configuration.V3VulnerabilityContainerName + "/");

            if (runMode == RunMode.Update && !await _storage.ExistsAsync(indexStorageUri.ToString(), CancellationToken.None))
            {
                _logger.LogWarning("Update mode was set, but {IndexPath} doesn't exist. Falling back to regeneration.", indexStorageUri.ToString());
                runMode = RunMode.Regenerate;
            }

            switch (runMode)
            {
                case RunMode.Update:
                    await RunUpdate(stringContentOutput, indexStorageUri, currentTime);
                    _telemetryService.TrackUpdateRun(vulnerabilityCount);
                    break;
                case RunMode.Regenerate:
                    await RunRegenerate(stringContentOutput, indexStorageUri, currentTime);
                    _telemetryService.TrackRegenerationRun(vulnerabilityCount);
                    break;
                case RunMode.None:
                default:
                    var exception = new InvalidOperationException();
                    _logger.LogError(exception, "RunMode was unable to be determined correctly.");
                    // Throw here. We don't want to break anything.
                    throw exception;
            }
        }

        public async Task<int> WriteVulnerabilitiesAsync(IEnumerable<Tuple<PackageVulnerability, bool>> vulnerabilities)
        {
            _firstVulnWrittenTimestamp = _firstVulnWrittenTimestamp == DateTimeOffset.MinValue ? DateTimeOffset.UtcNow : _firstVulnWrittenTimestamp;
            foreach (var vulnerability in vulnerabilities)
            {
                await WriteVulnerabilityAsync(vulnerability.Item1, vulnerability.Item2);
            }

            return vulnerabilities.Count();
        }

        public Task WriteVulnerabilityAsync(PackageVulnerability packageVulnerability, bool wasWithdrawn)
        {
            _firstVulnWrittenTimestamp = _firstVulnWrittenTimestamp == DateTimeOffset.MinValue ? DateTimeOffset.UtcNow : _firstVulnWrittenTimestamp;
            try
            {
                if (!wasWithdrawn)
                {
                    foreach (var vulnerableRange in packageVulnerability.AffectedRanges)
                    {
                        var idToWrite = vulnerableRange.PackageId.ToLowerInvariant();
                        var advisory = new Advisory
                        {
                            Url = packageVulnerability.AdvisoryUrl,
                            Severity = (int)packageVulnerability.Severity,
                            Versions = vulnerableRange.PackageVersionRange
                        };

                        if (!_vulnerabilityDict.ContainsKey(idToWrite))
                        {
                            _vulnerabilityDict.Add(idToWrite, new List<Advisory>());
                        }

                        _vulnerabilityDict[idToWrite].Add(advisory);
                    }
                }
                else
                {
                    foreach (var vulnerableRange in packageVulnerability.AffectedRanges)
                    {
                        var idToUpdate = vulnerableRange.PackageId.ToLowerInvariant();
                        if (_vulnerabilityDict.ContainsKey(idToUpdate))
                        {
                            var currentAdvisories = _vulnerabilityDict[idToUpdate];
                            var advisoriesToRemove = new List<Advisory>();
                            foreach (var advisory in currentAdvisories) 
                            { 
                                if(advisory.Url.Equals(packageVulnerability.AdvisoryUrl))
                                {
                                    advisoriesToRemove.Add(advisory);
                                }
                            }

                            foreach (var advisory in advisoriesToRemove)
                            {
                                _vulnerabilityDict[idToUpdate].Remove(advisory);
                            }
                        }
                    }
                }
            }
            catch (Exception e)
            {
                _logger.LogError(e, "WriteVulnerability Failed for Advisory Url {AdvisoryUrl}", packageVulnerability.AdvisoryUrl);
                throw;
            }

            return Task.CompletedTask;
        }

        public async Task<RunMode> DetermineRunMode(ReadWriteCursor<DateTimeOffset> cursor)
        {
            var mode = RunMode.Update;
            await cursor.Load(CancellationToken.None);
            if (DateTimeOffset.Compare(cursor.Value.AddDays(_configuration.DaysBeforeBaseStale), DateTimeOffset.Now) <= 0)
            {
                mode = RunMode.Regenerate;
            }
            return mode;
        }

        public async Task<bool> ShouldRegenerateForSpecialCase(Dictionary<string,List<Advisory>> baseContent)
        {
            /// Currently, if client sees the same vulnerability URL in update and base, it will be displayed twice.
            ///   (Unless everything from version range to package id to vulnerability URL is the same. In this case only, client can dedupe)
            /// This is undesirable behavior, but we are tracking to see how much this actually happens.
            /// We will load the current base.json and compare URLs to the update to see if there are repeats
            /// Ideally, once this is resolved, we can simply remove this method and calls in a single piece.
            /// if this method determines we should regenerate, it will update the cursor to Unix epoch and end the run.
            /// This way, the next run will pick up regeneration.
            _logger.LogInformation("Checking for special case trigger.");
            var baseUriHashSet = new HashSet<string>(StringComparer.OrdinalIgnoreCase);

            foreach(var entry in baseContent)
            {
                foreach(var advisory in entry.Value)
                {
                    baseUriHashSet.Add(advisory.Url);
                }
            }

            foreach(var entry in _vulnerabilityDict)
            {
                foreach(var advisory in entry.Value)
                {
                    if (!baseUriHashSet.Add(advisory.Url))
                    {
                        _cursor.Value = DateTimeOffset.MinValue;
                        await _cursor.Save(CancellationToken.None);
                        _logger.LogInformation("Special case triggered for advisory url {AdvisoryUrl}", advisory.Url);
                        return true;
                    }
                }
            }

            return false;
        }

        private async Task RunUpdate(string stringContentOutput, Uri indexStorageUri, string currentTime)
        {
            _logger.LogInformation("Beginning Update Run");
            var basePath = "";
            UriBuilder updateUriBuilder = new UriBuilder(_configuration.V3BaseUrl);
            StringBuilder updateUriPathBuilder = new StringBuilder();
            updateUriPathBuilder.Append(_configuration.V3VulnerabilityContainerName + "/");

            // Load the index and....
            var indexStorageContent = await _storage.LoadString(indexStorageUri, CancellationToken.None);
            var indexEntries = JsonConvert.DeserializeObject<IndexEntry[]>(indexStorageContent);

            // Scan the index for the base file so we can determine which folder we need to write in for the update.
            foreach (var entry in indexEntries)
            {
                if (entry.Name.Equals(BASE_ENTRY_NAME))
                {
                    basePath = entry.Id;
                    break;
                }
            }

            var pathParts = basePath.Split('/');
            // the folder should always be 2 entries in from the end of the list
            // Ex: https://api.nuget.org/v3/vulnerabilities/1234.56.78.90.12.34/vulnerability.base.json
            //                                                   HERE
            var baseFolder = pathParts[pathParts.Length - 2];

            // Start special case block
            var currentBaseContentUri = _storage.ResolveUri(Path.Combine(baseFolder, _configuration.BaseFileName));
            var currentBaseContent = await _storage.LoadString(currentBaseContentUri, CancellationToken.None);
            var baseContentObject = JsonConvert.DeserializeObject<Dictionary<string, List<Advisory>>>(currentBaseContent);

            if (await ShouldRegenerateForSpecialCase(baseContentObject))
            {
                _telemetryService.TrackSpecialCaseTrigger();
                return;
            }
            // End special case block

            updateUriPathBuilder.Append(baseFolder + "/");
            updateUriPathBuilder.Append(currentTime + "/");
            updateUriPathBuilder.Append(_configuration.UpdateFileName);
            updateUriBuilder.Path = updateUriPathBuilder.ToString();

            var updateStorageUri = _storage.ResolveUri(Path.Combine(baseFolder, currentTime, _configuration.UpdateFileName));

            for (var i = 0; i < indexEntries.Length; i++)
            {
                if (indexEntries[i].Name.Equals(UPDATE_ENTRY_NAME))
                {
                    indexEntries[i].Id = updateUriBuilder.Uri.AbsoluteUri;
                    indexEntries[i].Updated = DateTime.UtcNow;
                    break;
                }
            }

            _logger.LogInformation("Writing update to files");
            var updateContent = new StringStorageContent(stringContentOutput, contentType: JsonContentType, cacheControl: _configuration.UpdateCacheControlHeader);
            var indexContent = new StringStorageContent(JsonConvert.SerializeObject(indexEntries), contentType: JsonContentType, cacheControl: _configuration.IndexCacheControlHeader);

            await _storage.Save(updateStorageUri, updateContent, true, CancellationToken.None);
            await _storage.Save(indexStorageUri, indexContent, true, CancellationToken.None);
        }

        private async Task RunRegenerate(string stringContentOutput, Uri indexStorageUri, string currentTime)
        {
            _logger.LogInformation("Begin Regeneration Run");
            var updatedTime = DateTime.UtcNow;
            UriBuilder baseUriBuilder = new UriBuilder(_configuration.V3BaseUrl);
            StringBuilder baseUriPathBuilder = new StringBuilder();
            baseUriPathBuilder.Append(_configuration.V3VulnerabilityContainerName + "/");
            baseUriPathBuilder.Append(currentTime + "/");
            baseUriPathBuilder.Append(_configuration.BaseFileName);
            baseUriBuilder.Path = baseUriPathBuilder.ToString();

            var baseStorageUri = _storage.ResolveUri(Path.Combine(currentTime, _configuration.BaseFileName));

            UriBuilder updateUriBuilder = new UriBuilder(_configuration.V3BaseUrl);
            StringBuilder updateUriPathBuilder = new StringBuilder();
            updateUriPathBuilder.Append(_configuration.V3VulnerabilityContainerName + "/");
            updateUriPathBuilder.Append(currentTime + "/");
            updateUriPathBuilder.Append(currentTime + "/");
            updateUriPathBuilder.Append(_configuration.UpdateFileName);
            updateUriBuilder.Path = updateUriPathBuilder.ToString();

            var updateStorageUri = _storage.ResolveUri(Path.Combine(currentTime, currentTime, _configuration.UpdateFileName));

            var indexEntries = new IndexEntry[] {
                new IndexEntry
                {
                    Name=BASE_ENTRY_NAME,
                    Id=baseUriBuilder.Uri.AbsoluteUri,
                    Updated=updatedTime,
                    Comment="The base data for vulnerability update periodically"
                },
                new IndexEntry
                {
                    Name=UPDATE_ENTRY_NAME,
                    Id=updateUriBuilder.Uri.AbsoluteUri,
                    Updated=updatedTime,
                    Comment="The patch data for the vulnerability. Contains all the vulnerabilities since base was last updated."
                },
            };

            _logger.LogInformation("Writing Regenerated files out");
            var baseContent = new StringStorageContent(stringContentOutput, contentType: JsonContentType, cacheControl: _configuration.BaseCacheControlHeader);
            var updateContent = new StringStorageContent("{}", contentType: JsonContentType, cacheControl: _configuration.UpdateCacheControlHeader);

            var indexContent = new StringStorageContent(JsonConvert.SerializeObject(indexEntries), contentType: JsonContentType, cacheControl: _configuration.IndexCacheControlHeader);

            await _storage.Save(baseStorageUri, baseContent, true, CancellationToken.None);
            await _storage.Save(updateStorageUri, updateContent, true, CancellationToken.None);
            await _storage.Save(indexStorageUri, indexContent, true, CancellationToken.None);

            _cursor.Value = _firstVulnWrittenTimestamp;
            await _cursor.Save(CancellationToken.None);
        }
    }
}
